package com.google.p2pbench;

import android.content.BroadcastReceiver;
import android.content.Context;
import android.content.Intent;
import android.content.IntentFilter;
import android.net.nsd.NsdManager;
import android.net.nsd.NsdServiceInfo;
import android.net.wifi.SupplicantState;
import android.net.wifi.WifiConfiguration;
import android.net.wifi.WifiInfo;
import android.net.wifi.WifiManager;
import android.net.wifi.WifiManager.WifiLock;
import android.os.Handler;
import android.os.Parcel;
import android.os.Parcelable;
import android.util.Log;
import android.util.Pair;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.Inet4Address;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.net.SocketException;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;

public class WifiAdhocConnectivity implements Connectivity {
  private static final String TAG = "WifiAdhocConnectivity";

  private static final InetAddress BROADCAST_ADDRESS;

  static {
    InetAddress address;
    try {
      address = InetAddress.getByName("169.254.255.255");
    } catch (java.io.IOException e) {
      address = null; // WTF?
    }
    BROADCAST_ADDRESS = address;
  }

  private static final int BROADCAST_PORT = 7948;

  private final Context mContext;
  private final WifiManager mWifiManager;
  private final WifiLock mLock;
  private boolean mInitiallyOn;
  private WifiConfiguration mWifiConfig;
  private int mNetworkId = -1;
  private WifiStateObserver mWifiStateObserver;

  private DatagramSocket mDatagramSocket;
  private InetAddress mMyAddress;

  private BetterNsdManager mNsdManager;

  private Listener mListener;
  private final Map<Pair<String, Service>, ServiceBackend> mServices =
      new HashMap<Pair<String, Service>, ServiceBackend>();
  private final Map<Pair<String, Discoverer>, DiscovererBackend> mDiscoverers =
      new HashMap<Pair<String, Discoverer>, DiscovererBackend>();
  private ServiceBackend mBroadcastService;

  private final Handler mHandler = new Handler();

  private class WifiStateObserver extends BroadcastReceiver {
    @Override
    public void onReceive(Context context, Intent intent) {
      mHandler.post(new Runnable() {
        @Override
        public void run() {
          Log.e(TAG, "WIFI_STATE_CHANGED_ACTION");
          onWifiChange();
        }
      });
    }
  }

  private void registerObserver() {
    if (mWifiStateObserver != null) return;
    mWifiStateObserver = new WifiStateObserver();
    IntentFilter filter = new IntentFilter();
    filter.addAction(WifiManager.WIFI_STATE_CHANGED_ACTION);
    filter.addAction(WifiManager.NETWORK_STATE_CHANGED_ACTION);
    mContext.registerReceiver(mWifiStateObserver, filter);
  }

  private void unregisterObserver() {
    if (mWifiStateObserver == null) return;
    Log.w(TAG, "Unregistering");
    mContext.unregisterReceiver(mWifiStateObserver);
    mWifiStateObserver = null;
  }

  private void onWifiChange() {
    if (mNetworkId < 0) {
      mNetworkId = mWifiManager.addNetwork(mWifiConfig);
      if (mNetworkId < 0)
        return;  // TODO: try again later...
    }
    mLock.acquire();

    WifiInfo wifiInfo = mWifiManager.getConnectionInfo();
    Log.w(TAG, "onWifiChange " + wifiInfo.getNetworkId() + wifiInfo.getSupplicantState());
    if (mNetworkId != wifiInfo.getNetworkId()) {
      Log.w(TAG, "enableNetwork " + mNetworkId);
      if (!mWifiManager.enableNetwork(mNetworkId, true /* disableOthers */)) {
        if (mListener != null)
          mListener.onError(new ConnectivityException("Could not enable network"));
      }
      return;
    }

    if (!wifiInfo.getSupplicantState().equals(SupplicantState.COMPLETED)) return;

    unregisterObserver();
    // NOTE: this is too often set to 0.0.0.0 if not connected initially.
    // mMyAddress = WifiAdhocApi.intToInetAddress(wifiInfo.getIpAddress());
    mMyAddress = Util.getIpAddress();
    Log.w(TAG, "myAddress " + mMyAddress);

    if (mDatagramSocket != null) return;

    try {
      mDatagramSocket = new DatagramSocket(null);
      mDatagramSocket.setReuseAddress(true);
      mDatagramSocket.bind(new InetSocketAddress(BROADCAST_PORT));
    } catch (SocketException e) {
      Log.e(TAG, "Could not create datagram socket", e);
      if (mListener != null)
        mListener.onError(new ConnectivityException("Could not bind broadcast socket", e));
      return;
    }

    // mNsdManager = (NsdManager) mContext.getSystemService(Context.NSD_SERVICE);
    try {
      mNsdManager = new BetterNsdManager(BROADCAST_ADDRESS);
    } catch (SocketException e) {
      if (mListener != null)
        mListener.onError(new ConnectivityException("Could not create NSD", e));
      return;
    }

    if (mListener != null) mListener.onConnected();
  }

  private static class PeerImpl implements Peer {
    final SocketAddress mAddress;

    PeerImpl(InetAddress address, int port) {
      mAddress = new java.net.InetSocketAddress(address, port);
    }

    PeerImpl(SocketAddress address) {
      mAddress = address;
    }

    @Override
    public String getDescription() {
      return "[WifiAdhoc " + mAddress.toString() + "]";
    }

    @Override
    public int describeContents() {
      return 0;
    }

    @SuppressWarnings("unused")
    public static final Parcelable.Creator<PeerImpl> CREATOR = new Parcelable.Creator<PeerImpl>() {
      @Override
      public PeerImpl createFromParcel(Parcel in) {
        return new PeerImpl((SocketAddress) in.readSerializable());
      }

      @Override
      public PeerImpl[] newArray(int size) {
        return new PeerImpl[size];
      }
    };

    @Override
    public void writeToParcel(Parcel dest, int flags) {
      dest.writeSerializable(mAddress);
    }
  }

  private static class SocketImpl implements Socket {
    java.net.Socket mSocket;

    SocketImpl(java.net.Socket socket) {
      mSocket = socket;
    }

    @Override
    public void close() {
      try {
        mSocket.close();
      } catch (java.io.IOException e) {
        // ignore
      }
    }

    @Override
    public boolean isConnected() {
      return mSocket.isConnected();
    }

    @Override
    public InputStream getInputStream() throws ConnectivityException {
      try {
        return mSocket.getInputStream();
      } catch (java.io.IOException e) {
        throw new ConnectivityException("Socket failed", e);
      }
    }

    @Override
    public OutputStream getOutputStream() throws ConnectivityException {
      try {
        return mSocket.getOutputStream();
      } catch (java.io.IOException e) {
        throw new ConnectivityException("Socket failed", e);
      }
    }
  }

  private class ServiceBackend {
    final String mServiceType;
    final Service mService;

    java.net.ServerSocket mServerSocket;
    java.net.DatagramSocket mMessageSocket;
    Thread mAcceptorThread;
    Thread mDatagramThread;
    boolean mStopped = false;

    NsdManager.RegistrationListener mRegistrationListener =
        new NsdManager.RegistrationListener() {
      @Override
      public void onRegistrationFailed(NsdServiceInfo serviceInfo, int errorCode) {
        mRegistrationListener = null;
        mService.onError(new ConnectivityException("Registration failed " + errorCode));
      }

      @Override
      public void onServiceRegistered(NsdServiceInfo serviceInfo) {}

      @Override
      public void onServiceUnregistered(NsdServiceInfo serviceInfo) {}

      @Override
      public void onUnregistrationFailed(NsdServiceInfo serviceInfo, int errorCode) {}
    };

    Runnable mAcceptor = new Runnable() {
      @Override
      public void run() {
        try {
          while (true) {
            java.net.Socket socket = mServerSocket.accept();
            mService.onStream(new PeerImpl(socket.getRemoteSocketAddress()),
                new SocketImpl(socket));
          }
        } catch (IOException e) {
          mService.onError(new ConnectivityException("Acceptor failed", e));
          shutdown();
        }
      }
    };

    Runnable mMessageReader = new Runnable() {
      @Override
      public void run() {
        try {
          byte[] data = new byte[2500];
          DatagramPacket packet = new DatagramPacket(data, data.length);
          DatagramSocket socket = mMessageSocket;
          if (socket == null) return;
          while (!mStopped) {
            socket.receive(packet);
            ByteBuffer msg = ByteBuffer.wrap(packet.getData(), packet.getOffset(),
                packet.getLength());
            mService.onMessage(new PeerImpl(packet.getSocketAddress()), msg);
          }
        } catch (IOException e) {
          if (!mStopped) mService.onError(new ConnectivityException("DatagramReader failed", e));
        }
      }
    };

    // Broadcast service
    ServiceBackend(Service service) throws ConnectivityException {
      mServiceType = "BROADCAST";
      mService = service;
      try {
        mMessageSocket = new DatagramSocket();
        mDatagramSocket.setReuseAddress(true);
        mDatagramSocket.bind(new InetSocketAddress(BROADCAST_PORT));
      } catch (IOException e) {
        throw new ConnectivityException("Failed to start broadcast service", e);
      }

      mDatagramThread = new Thread(mMessageReader, mService.toString() + " datagram reader");
      mDatagramThread.start();
    }

    ServiceBackend(String serviceType, Service service) throws ConnectivityException {
      assert mRegistrationListener != null;
      mServiceType = serviceType;
      mService = service;

      for (int i = 0; i < 4; ++i) {
        try {
          // One port for both UDP and TCP so that there's only one DNS-SD announcement.
          // TODO(szym): DNS-SD announcement should be per transport (_tcp or _udp).
          int port = (new Random()).nextInt(1024) + (65536 - 1024);
          mMessageSocket = new DatagramSocket(port);
          mServerSocket = new java.net.ServerSocket(port);
        } catch (IOException e) {
          // try again
          if (mServerSocket != null) {
            try {
              mServerSocket.close();
            } catch (IOException f) {
              // ignore
            }
            mServerSocket = null;
          }
          if (mMessageSocket != null) {
            mMessageSocket.close();
            mMessageSocket = null;
          }
        }
      }

      if (mServerSocket == null) {
        throw new ConnectivityException("Failed to start service");
      }

      mAcceptorThread = new Thread(mAcceptor, mService.toString() + " acceptor");
      mAcceptorThread.start();
      mDatagramThread = new Thread(mMessageReader, mService.toString() + " datagram reader");
      mDatagramThread.start();

      NsdServiceInfo info = new NsdServiceInfo();
      info.setServiceName(mServiceType); // TODO: Lame!
      info.setServiceType(mServiceType);
      // NOTE: DOES NOT MATTER because system/netd/MDnsSdListener.cpp:429 uses NULL anyway!
      info.setHost(mMyAddress);
      info.setPort(mServerSocket.getLocalPort());
      Log.w(TAG, "register NSD " + info);
      mNsdManager.registerService(info, NsdManager.PROTOCOL_DNS_SD, mRegistrationListener);
    }

    PeerImpl getPeer() {
      return new PeerImpl(mServerSocket.getLocalSocketAddress());
    }

    void shutdown() {
      mStopped = true;
      if (mServerSocket != null) {
        Log.w(TAG, "unregister NSD");
        if (mRegistrationListener != null) {
          if (mNsdManager != null) mNsdManager.unregisterService(mRegistrationListener);
          mRegistrationListener = null;
        }
        try {
          mServerSocket.close();
        } catch (IOException e) {
          // ignore
        }
        mServerSocket = null;
      }
      if (mMessageSocket != null) {
        mMessageSocket.close();
        mMessageSocket = null;
      }
    }
  }

  private class DiscovererBackend {
    String mServiceType;
    Discoverer mDiscoverer;

    NsdManager.DiscoveryListener mDiscoveryListener = new NsdManager.DiscoveryListener() {
      @Override
      public void onDiscoveryStarted(String serviceType) {
      }

      @Override
      public void onDiscoveryStopped(String serviceType) {
      }

      @Override
      public void onServiceFound(NsdServiceInfo serviceInfo) {
        // TODO(szym): provide a way to indicate it's on the same host?
        Log.w(TAG, "onServiceFound " + serviceInfo);
        Log.w(TAG, "check " + serviceInfo.getServiceType() + " " + mServiceType + " "
            + mServiceType.equals(serviceInfo.getServiceType()));
        if (!mServiceType.equals(serviceInfo.getServiceType())) return; // Ignore unknown service.
        Log.w(TAG, "Resolving!");
        mNsdManager.resolveService(serviceInfo, mResolveListener);
      }

      @Override
      public void onServiceLost(NsdServiceInfo serviceInfo) {
      }

      @Override
      public void onStartDiscoveryFailed(String serviceType, int errorCode) {
        mDiscoveryListener = null;
      }

      @Override
      public void onStopDiscoveryFailed(String serviceType, int errorCode) {
      }
    };

    NsdManager.ResolveListener mResolveListener = new NsdManager.ResolveListener() {
      @Override
      public void onResolveFailed(NsdServiceInfo serviceInfo, int errorCode) {
        // Ignore.
        Log.e(TAG, "Resolve failed " + errorCode);
      }

      @Override
      public void onServiceResolved(NsdServiceInfo serviceInfo) {
        Log.w(TAG, "Resolved! " + serviceInfo);
        // BroadcastReceiver might be slower than the socket, so let's fetch the
        // address on demand (hopefully benign data race).
        if (mMyAddress == null) mMyAddress = Util.getIpAddress();
        if (mMyAddress.equals(serviceInfo.getHost())) return; // Ignore self.
        Peer peer = new PeerImpl(serviceInfo.getHost(), serviceInfo.getPort());
        mDiscoverer.onDiscovery(serviceInfo.getServiceType(), peer);
      }
    };

    DiscovererBackend(String serviceType, Discoverer discoverer) {
      mServiceType = serviceType;
      mDiscoverer = discoverer;

      mNsdManager.discoverServices(mServiceType, NsdManager.PROTOCOL_DNS_SD, mDiscoveryListener);
    }

    void shutdown() {
      if (mDiscoveryListener != null) {
        if (mNsdManager != null) mNsdManager.stopServiceDiscovery(mDiscoveryListener);
        mDiscoveryListener = null;
      }
    }
  }


  public WifiAdhocConnectivity(Context context) {
    mContext = context;
    mWifiManager = (WifiManager) mContext.getSystemService(Context.WIFI_SERVICE);
    // For IBSS association, it has to be HIGH_PERF to have any effect.
    mLock = mWifiManager.createWifiLock(WifiManager.WIFI_MODE_FULL_HIGH_PERF, TAG);
  }

  @Override
  public boolean supportsMessages() {
    return true;
  }

  @Override
  public boolean connect(Listener listener) {
    Util.assertOnThread(mHandler);
    mListener = listener;

    WifiConfiguration config = new WifiConfiguration();
    config.SSID = WifiAdhocApi.convertToSSID("AdhocBenchmark");
    // Disable encryption.
    config.allowedKeyManagement.clear();
    config.allowedKeyManagement.set(WifiConfiguration.KeyMgmt.NONE);

    // FIXME: not actually needed / used.
    InetAddress dnsAddress;
    try {
      dnsAddress = Inet4Address.getByAddress(new byte[] {224 - 256, 0, 0, 251 - 256});
    } catch (UnknownHostException e) {
      // IMPOSSIBRU!
      return false;
    }
    InetAddress myAddress;
    try {
      byte[] myAddressBytes = new byte[4];
      (new Random()).nextBytes(myAddressBytes);
      myAddressBytes[0] = 169 - 256;
      myAddressBytes[1] = 254 - 256;
      myAddress = Inet4Address.getByAddress(myAddressBytes);
    } catch (UnknownHostException e) {
      // IMPOSSIBRU!
      return false;
    }
    if (!WifiAdhocApi.configure(config, myAddress, 16, dnsAddress)) return false;
    mWifiConfig = config;
    registerObserver();

    mInitiallyOn = mWifiManager.isWifiEnabled();
    if (mInitiallyOn) {
      onWifiChange();
      return true;
    }
    return mWifiManager.setWifiEnabled(true);
  }

  @Override
  public void shutdown(int mode) {
    Util.assertOnThread(mHandler);
    if (mNetworkId < 0) return;
    if (mLock.isHeld()) mLock.release();
    mWifiManager.disconnect();
    mWifiManager.removeNetwork(mNetworkId);
    mNetworkId = -1;
    if (mDatagramSocket != null)
      mDatagramSocket.close(); // this kills the Receiver thread
    mDatagramSocket = null;
    if (mNsdManager != null) mNsdManager.shutdown();
    mNsdManager = null;
    unregisterObserver();
    try {
      setBroadcastService(null);
    } catch (ConnectivityException e) {
    }
    for (ServiceBackend b : mServices.values()) {
      b.shutdown();
    }
    for (DiscovererBackend b : mDiscoverers.values()) {
      b.shutdown();
    }
    if (mode == SHUTDOWN_OFF || (mode == SHUTDOWN_PREVIOUS && !mInitiallyOn)) {
      mWifiManager.setWifiEnabled(false);
    } else {
      // lame API
      List<WifiConfiguration> networks = mWifiManager.getConfiguredNetworks();
      if (networks != null) {
        for (WifiConfiguration wc : networks) {
          mWifiManager.enableNetwork(wc.networkId, false);
        }
      }
      // mWifiManager.reconnect();
    }
  }

  @Override
  public Peer getBroadcastPeer() {
    return new PeerImpl(BROADCAST_ADDRESS, BROADCAST_PORT);
  }

  @Override
  public void setBroadcastService(Service service) throws ConnectivityException {
    Util.assertOnThread(mHandler);
    if (mBroadcastService != null)
      mBroadcastService.shutdown();
    mBroadcastService = null;
    if (service != null) mBroadcastService = new ServiceBackend(service);
  }

  @Override
  public void sendMessage(Peer peer, ByteBuffer msg) throws ConnectivityException {
    DatagramPacket packet = new DatagramPacket(msg.array(), msg.arrayOffset(), msg.remaining());
    packet.setSocketAddress(((PeerImpl) peer).mAddress);
    try {
      mDatagramSocket.send(packet);
    } catch (IOException e) {
      throw new ConnectivityException("Failed to send datagram", e);
    }
  }

  @Override
  public Socket getSocket(Peer peer) throws ConnectivityException {
    // NOTE: peer cannot be the broadcast peer!
    java.net.Socket socket = new java.net.Socket();
    try {
      socket.connect(((PeerImpl) peer).mAddress);
    } catch (IOException e) {
      try {
        socket.close();
      } catch (IOException f) {
        // ignore
      }
      throw new ConnectivityException("Could not connect", e);
    }
    return new SocketImpl(socket);
  }

  @Override
  public Peer announce(String serviceType, Service service) throws ConnectivityException {
    Util.assertOnThread(mHandler);
    ServiceBackend backend = new ServiceBackend(serviceType, service);
    mServices.put(Pair.create(serviceType, service), backend);
    return backend.getPeer();
  }

  @Override
  public void unnannounce(String serviceType, Service service) {
    Util.assertOnThread(mHandler);
    ServiceBackend backend = mServices.remove(Pair.create(serviceType, service));
    if (backend != null) backend.shutdown();
  }

  @Override
  public void discover(String serviceType, Discoverer listener) {
    Util.assertOnThread(mHandler);
    DiscovererBackend backend = new DiscovererBackend(serviceType, listener);
    mDiscoverers.put(Pair.create(serviceType, listener), backend);
  }

  @Override
  public void stopDiscovery(String serviceType, Discoverer listener) {
    Util.assertOnThread(mHandler);
    DiscovererBackend backend = mDiscoverers.remove(Pair.create(serviceType, listener));
    if (backend != null) backend.shutdown();
  }
}
